import asyncio
import os
import json
import sys
from typing import Optional
from  contextlib import AsyncExitStack
from openai.types.chat import ChatCompletionToolParam
from openai import OpenAI
from dotenv import load_dotenv

from mcp import ClientSession,StdioServerParameters
from mcp.client.stdio import stdio_client

from pypinyin import lazy_pinyin, Style

# 加载env文件，确保配置正确
load_dotenv()

class MCPClient:
    def __init__(self):
        """初始化MCP客户端"""
        self.write = None
        self.stdio = None
        self.exit_stack = AsyncExitStack()
        self.base_url=os.getenv("BASE_URL")
        self.model = os.getenv("MODEL")
        self.openai_api_key = os.getenv("OPENAI_API_KEY")

        if not self.openai_api_key:
            raise ValueError("OPENAI_API_KEY未设置")

        self.client=OpenAI(api_key=self.openai_api_key, base_url=self.base_url)
        self.session: Optional[ClientSession] = None
        self.exit_stack=AsyncExitStack()


    async def connect_to_server(self, server_script_path:str):
        """连接MCP服务器并列出可用工具"""
        is_python=server_script_path.endswith(".py")
        is_js=server_script_path.endswith(".js")
        if not (is_python or is_js):
            raise ValueError("服务器脚本必须以.py或.js结尾")

        command="python" if is_python else "node"
        server_params=StdioServerParameters(
            command=command,
            args=[server_script_path],
            env=None)
        # 启动MCP服务器并建立通信
        stdio_transport=await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.stdio,self.write=stdio_transport
        self.session= await self.exit_stack.enter_async_context(ClientSession(self.stdio,self.write))

        await self.session.initialize()

        # 列出MCP服务器上的工具
        response=await self.session.list_tools()
        tools=response.tools
        print("\n已连接到服务器，支持以下工具：",[tool.name for tool in tools])


    async def process_query(self, query:str)-> str:
        """
        使用大模型处理查询并调用可用的MCP工具（Function Calling）
        """
        messages=[{"role": "user","content": query}]

        response=await self.session.list_tools()

        available_tools = [
            ChatCompletionToolParam(
                type="function",
                function={
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.inputSchema
                }
            )
            for tool in response.tools
        ]

        response= self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            tools=available_tools
        )

        # 处理返回的内容
        content=response.choices[0]
        # 检查是否使用了工具
        if content.finish_reason == "tool_calls":
            # 何时需要使用工具就解析工具
            tool_call = content.message.tool_calls[0]
            tool_name=tool_call.function.name
            tool_args=json.loads(tool_call.function.arguments)

            # 如果调用的是 query_weather 工具，处理城市名称
            if tool_name == "query_weather" and "city" in tool_args:
                city = tool_args["city"]
                # 简单判断是否为中文城市名
                if any('\u4e00' <= c <= '\u9fff' for c in city):
                    # 转换为拼音，首字母大写
                    pinyin_city = ''.join([word.capitalize() for word in lazy_pinyin(city)])
                    tool_args["city"] = pinyin_city

            # 执行工具
            result=await self.session.call_tool(tool_name, tool_args)
            print(f"\n\n[Calling Tool: {tool_name} with args: {tool_args}]")

            # 将工具调用和结果添加到消息历史中
            messages.append(content.message.model_dump())
            messages.append({
                "role": "tool",
                "content": result.content[0].text,
                "tool_call_id": tool_call.id
            })

            # 将上面的结果再返回给大模型用于最终的效果
            response=self.client.chat.completions.create(model=self.model,messages=messages)
            return response.choices[0].message.content
        # 如果调用工具直接返回结果
        return content.message.content

    async def chat_loop(self):
        """运行交互式聊天循环"""
        print("\nMCP客户端已启动！输入'quit'退出")
        while True:
            try:
                query=input("\n you:").strip()
                if query.lower() == "quit":
                    break
                response=await self.process_query(query)
                print(f"\n ai: {response}")
            except Exception as e:
                print(f"\n Error: {e}")
    async def cleanup(self):
        """清理资源"""
        await self.exit_stack.aclose()

async def main():
    if len(sys.argv)<2:
        print("Usage: python client.py <server_address>")
        sys.exit(1)
    client=MCPClient()
    try:
        await client.connect_to_server(sys.argv[1])
        await client.chat_loop()
    finally:
        await client.cleanup()

if __name__ == "__main__":
    asyncio.run(main())